#include "stdafx.h"
#include "idbgserv.h"

#include "eventscb.h"
#include "out.h"

IDebugClient5* g_Client = NULL;
IDebugControl4* g_Control = NULL;
IDebugDataSpaces4* g_Data = NULL;
IDebugRegisters2* g_Registers = NULL;
IDebugSymbols3* g_Symbols = NULL;
IDebugSystemObjects4* g_SystemObjects = NULL;

StdioOutputCallbacks g_OutputCb;
EventCallbacks g_EventCb;

bool gbKernel = false;

void
StrFree (
	PStrList pCommand
	)
{
	if ( pCommand )
	{
		if ( pCommand->m_String )
			HeapFree( GetProcessHeap(), 0, pCommand->m_String );
		HeapFree( GetProcessHeap(), 0, pCommand );
	}
}

HRESULT
CreateInterfaces (
	)
{
	HRESULT hResult = DebugCreate( __uuidof( IDebugClient5 ), (void**) &g_Client );

	if ( !SUCCEEDED( hResult ) )
		return hResult;

	hResult |= g_Client->QueryInterface( __uuidof( IDebugControl4 ), (void**) &g_Control );
	hResult |= g_Client->QueryInterface( __uuidof( IDebugDataSpaces4 ), (void**) &g_Data );
	hResult |= g_Client->QueryInterface( __uuidof( IDebugRegisters2 ), (void**) &g_Registers );
	hResult |= g_Client->QueryInterface( __uuidof(IDebugSymbols3), (void**) &g_Symbols );
    hResult |= g_Client->QueryInterface( __uuidof(IDebugSystemObjects4), (void**) &g_SystemObjects );

	if ( SUCCEEDED( hResult ) )
		hResult = g_Client->SetEventCallbacks( &g_EventCb );
	if ( SUCCEEDED( hResult ) )
		hResult = g_Client->SetOutputCallbacks( &g_OutputCb );

	return hResult;
}

void
Cleanup (
	)
{
	if ( g_Control )
		g_Control->Release();

	if ( g_Data )
		g_Data->Release();

	if ( g_Registers )
		g_Registers->Release();

	if ( g_Symbols )
		g_Symbols->Release();

    if ( g_SystemObjects )
        g_SystemObjects->Release();
    
	if ( g_Client )
	{
		g_Client->EndSession( DEBUG_END_PASSIVE );
		g_Client->Release();
	}
}

DWORD
WINAPI
KernelConnectThread (
    LPVOID lpThreadParameter
    )
{
    gbKernel = true;
	CDbgService* pDbgService = (CDbgService*) lpThreadParameter;

	HRESULT hResult = CreateInterfaces();
	if ( SUCCEEDED( hResult) )
	{
		PStrList pCommand = pDbgService->Command_Get();
		if ( pCommand )
		{
			hResult = g_Client->AttachKernel( DEBUG_ATTACH_DEFAULT, pCommand->m_String );
			StrFree( pCommand );

			if ( SUCCEEDED( hResult) )
				pDbgService->EventLoop();
		}
	}

	return 0;
}

DWORD
WINAPI
AttachProcessThread (
    LPVOID lpThreadParameter
    )
{
    gbKernel = false;
	CDbgService* pDbgService = (CDbgService*) lpThreadParameter;

	HRESULT hResult = CreateInterfaces();
	if ( SUCCEEDED( hResult) )
	{
		hResult = g_Client->AttachProcess( 0, pDbgService->m_ProcessId, DEBUG_ATTACH_DEFAULT );

		if ( SUCCEEDED( hResult) )
			pDbgService->EventLoop();
	}

	return 0;
}

CDbgService::CDbgService()
{
	m_RefCount = 0;
	InitializeCriticalSection( &m_CommandCS );
	m_ListFirst = NULL;
	m_ListLast = NULL;
	m_hEvent = CreateEvent( NULL, FALSE, FALSE, NULL );
	m_hThread = NULL;

	m_ProcessId = 0;
};

CDbgService::~CDbgService()
{
	CloseHandle( m_hEvent );
	WaitForSingleObject( m_hThread, INFINITE );
	CloseHandle( m_hThread );
	DeleteCriticalSection( &m_CommandCS );

};

PStrList
CDbgService::Command_Get (
	)
{
	PStrList pResult = NULL;
	EnterCriticalSection( &m_CommandCS );

	if ( m_ListFirst )
	{
		pResult = m_ListFirst;
		m_ListFirst = m_ListFirst->m_pNext;

		if ( !m_ListFirst )
			m_ListLast = NULL;
	}

	LeaveCriticalSection( &m_CommandCS );
	
	return pResult;
}

void
CDbgService::Command_Push (
	ULONG Flags,
	ULONG OutputControl,
	BSTR bString
	)
{
	PCHAR pchString;

	int result = WideCharToMultiByte (
		CP_ACP,
		0,
		bString,
		SysStringByteLen(bString) / sizeof(WCHAR),
		NULL,
		0,
		NULL,
		NULL
		);

	if ( !result )
		return;

	PStrList pCommand = (PStrList) HeapAlloc( GetProcessHeap(), 0, sizeof(StrList) );
	pchString = (PCHAR) HeapAlloc( GetProcessHeap(), 0, result + sizeof(CHAR) );
	if ( pCommand && pchString )
	{
		result = WideCharToMultiByte (
			CP_ACP,
			0,
			bString,
			SysStringByteLen(bString) / sizeof(WCHAR),
			pchString,
			result,
			NULL,
			NULL
			);

		pchString[result] = 0;
		pCommand->m_Flags = Flags;
		pCommand->m_OutputControl = OutputControl;
		pCommand->m_String = pchString;
		pCommand->m_pNext = NULL;

		EnterCriticalSection( &m_CommandCS );

		if ( !m_ListLast )
		{
			m_ListFirst = pCommand;
			m_ListLast = pCommand;
		}
		else
		{
			m_ListLast->m_pNext = pCommand;
			m_ListLast = pCommand;
		}

		LeaveCriticalSection( &m_CommandCS );
		SetEvent( m_hEvent );

		return;
	}

	if ( pCommand )
		HeapFree( GetProcessHeap(), 0, pCommand );

	if ( pchString )
		HeapFree( GetProcessHeap(), 0, pchString );
}

void
CDbgService::EventLoop (
	)
{
	HRESULT hResult = S_OK;
	
	while( SUCCEEDED( hResult ) )
	{
		ULONG ExecStatus;

		hResult = g_Control->WaitForEvent( DEBUG_WAIT_DEFAULT, INFINITE );
		if ( !SUCCEEDED( hResult ) )
		{
			hResult = g_Control->GetExecutionStatus( &ExecStatus );

			if ( S_OK == hResult && DEBUG_STATUS_NO_DEBUGGEE == ExecStatus )
				break; // The session ended so we can quit.

			return; // There was a real error.
		}

		hResult = g_Control->GetExecutionStatus( &ExecStatus );
		while ( SUCCEEDED( hResult ) && DEBUG_STATUS_BREAK == ExecStatus)
		{
			// command loop
			WaitForSingleObject( m_hEvent, INFINITE );

			// execute single command
			PStrList pCommand = Command_Get();
			if ( !pCommand )
				break;
				
			g_Control->Execute( pCommand->m_OutputControl, pCommand->m_String,  pCommand->m_Flags );
			StrFree( pCommand );
			hResult = g_Control->GetExecutionStatus( &ExecStatus );
		}
	}
}

BSTR
CDbgService::CreateBstr (
    __in PSTR pstrString
    )
{
    BSTR bstrRet = NULL;
	OLECHAR* psz = NULL;

	int result = MultiByteToWideChar (
		CP_ACP,
		0,
		pstrString,
		lstrlenA( pstrString ),
		NULL,
		0
		);

	if ( result )
	{
		psz = (OLECHAR*) HeapAlloc( GetProcessHeap(), 0, ( result + 1 ) * sizeof( WCHAR ) );
		if ( psz )
		{
			result = MultiByteToWideChar (
				CP_ACP,
				0,
				pstrString,
				lstrlenA( pstrString ),
				psz,
				result
				);
			
			psz[result] = 0;

			if ( !result )
			{
				HeapFree( GetProcessHeap(), 0, psz );
				psz = NULL;
			}
		}
	}

	if ( psz )
	{
		bstrRet = SysAllocString( psz );

		HeapFree( GetProcessHeap(), 0, psz );
	}

    return bstrRet;
}

// interface

HRESULT
CDbgService::QueryInterface (
	REFIID riid,
	void **ppvObject
	)
{
	*ppvObject = NULL;

	if (IsEqualIID(riid, __uuidof(IUnknown)) ||
		IsEqualIID(riid, __uuidof(IDbgService)))
	{
		*ppvObject = this;
		AddRef( );
		return S_OK;
	}

	return E_NOINTERFACE;
}

ULONG
CDbgService::AddRef (
	)
{
	m_RefCount++;
	return m_RefCount;
}

ULONG
CDbgService::Release (
	)
{
	m_RefCount--;
	return m_RefCount;
}

HRESULT
CDbgService::GetTypeInfoCount (
	UINT *pctinfo
	)
{
	UNREFERENCED_PARAMETER( pctinfo );
	OutputDebugString( L"GetTypeInfoCount\n" );

	return E_NOTIMPL;
}

HRESULT
CDbgService::GetTypeInfo (
	UINT iTInfo,
	LCID lcid,
	ITypeInfo **ppTInfo
	)
{
	UNREFERENCED_PARAMETER( iTInfo );
	UNREFERENCED_PARAMETER( lcid );
	UNREFERENCED_PARAMETER( ppTInfo );

	OutputDebugString( L"GetTypeInfo\n" );
	return E_NOTIMPL;
}

HRESULT
CDbgService::GetIDsOfNames (
	REFIID riid,
	LPOLESTR *rgszNames,
	UINT cNames,
	LCID lcid,
	DISPID *rgDispId
	)
{
	UNREFERENCED_PARAMETER( riid );
	UNREFERENCED_PARAMETER( rgszNames );
	UNREFERENCED_PARAMETER( cNames );
	UNREFERENCED_PARAMETER( lcid );
	UNREFERENCED_PARAMETER( rgDispId );

	OutputDebugString( L"GetIDsOfNames\n" );
	return E_NOTIMPL;
}

HRESULT
CDbgService::Invoke (
	DISPID dispIdMember,
	REFIID riid,
	LCID lcid,
	WORD wFlags,
	DISPPARAMS *pDispParams,
	VARIANT *pVarResult,
	EXCEPINFO *pExcepInfo,
	UINT *puArgErr
	)
{
	UNREFERENCED_PARAMETER( dispIdMember );
	UNREFERENCED_PARAMETER( riid );
	UNREFERENCED_PARAMETER( lcid );
	UNREFERENCED_PARAMETER( wFlags );
	UNREFERENCED_PARAMETER( pDispParams );
	UNREFERENCED_PARAMETER( pVarResult );
	UNREFERENCED_PARAMETER( pExcepInfo );
	UNREFERENCED_PARAMETER( puArgErr );

	OutputDebugString( L"Invoke\n" );
	return E_NOTIMPL;
}

HRESULT
CDbgService::KernelConnect ( 
	__in BSTR pVal
	)
{
	Command_Push( 0, 0, pVal );
	SysFreeString( pVal );
	m_hThread = CreateThread( NULL, 0, KernelConnectThread, this, 0, NULL );

	return S_OK;
}

HRESULT
CDbgService::AttachProcess ( 
	__in ULONG ProcessId
	)
{
	m_ProcessId = ProcessId;
	m_hThread = CreateThread( NULL, 0, AttachProcessThread, this, 0, NULL );

	return S_OK;
}

HRESULT
CDbgService::IsPointer64Bit (
    )
{
    HRESULT hResult = g_Control->IsPointer64Bit();
    if ( S_OK == hResult )
        return S_OK;

    return E_FAIL;
}

HRESULT
CDbgService::GetOutput (
	__out BSTR* ppVal
	)
{
	PStrList pOutput = g_OutputCb.GetOutput();
	*ppVal = NULL;

	if ( pOutput )
	{
        *ppVal = CreateBstr( pOutput->m_String );
		StrFree( pOutput );
	}

	return S_OK;
}

HRESULT
CDbgService::Break (
	)
{
	g_Control->SetInterrupt( DEBUG_INTERRUPT_ACTIVE );
	
	return S_OK;
}

HRESULT
CDbgService::Execute (
	__in BSTR pCommand
	)
{
	Command_Push( 0, 0, pCommand );
	SysFreeString( pCommand );

	return S_OK;
}

HRESULT
CDbgService::GetExecutionStatus (
	__out PULONG pStatus
	)
{
	return g_Control->GetExecutionStatus( pStatus );
}

HRESULT
CDbgService::GetStackFrameCount (
	__out ULONG* pFrameCount
	)
{
	DEBUG_STACK_FRAME Stacks[100];

	HRESULT hResult = g_Control->GetStackTrace( 0, 0, 0, &Stacks[0], _countof(Stacks), pFrameCount );
    if ( FAILED( hResult ) )
		*pFrameCount = 0;

	return S_OK;
}

HRESULT
CDbgService::GetStackFrame (
	__in ULONG Frame,
	__out DBG_STACK_FRAME* pFrame
	)
{
	RtlZeroMemory( pFrame, sizeof(DBG_STACK_FRAME) );

	ULONG Frames = 0;
	DEBUG_STACK_FRAME Stacks[100];

	HRESULT hResult = g_Control->GetStackTrace( 0, 0, 0, &Stacks[0], _countof(Stacks), &Frames );
    
	if ( FAILED( hResult ) )
		Frames = 0;

	if ( Frames < Frame )
		return E_INVALIDARG;

	g_Symbols->SetScope( 0, &Stacks[Frame], NULL, 0 );

	pFrame->InstructionOffset = Stacks[Frame].InstructionOffset;
	pFrame->ReturnOffset = Stacks[Frame].ReturnOffset;
	pFrame->FrameOffset = Stacks[Frame].FrameOffset;
	pFrame->StackOffset = Stacks[Frame].StackOffset;
	pFrame->FuncTableEntry = Stacks[Frame].FuncTableEntry;
	pFrame->Params[0] = Stacks[Frame].Params[0];
	pFrame->Params[1] = Stacks[Frame].Params[1];
	pFrame->Params[2] = Stacks[Frame].Params[2];
	pFrame->Params[3] = Stacks[Frame].Params[3];
	pFrame->Virtual = Stacks[Frame].Virtual;
	pFrame->FrameNumber = Stacks[Frame].FrameNumber;

    /*	IDebugSymbolGroup2 *pDbgSymGroup = NULL;
    hResult = g_Symbols->GetScopeSymbolGroup2( DEBUG_SCOPE_GROUP_ALL, NULL, &pDbgSymGroup );
	if ( SUCCEEDED( hResult ) )
	{
		ULONG ulCount;
		hResult = pDbgSymGroup->GetNumberSymbols ( &ulCount );

		if ( SUCCEEDED( hResult ) && ulCount )
		{
			PDEBUG_SYMBOL_PARAMETERS pSymParams = new DEBUG_SYMBOL_PARAMETERS[ ulCount ] ;
			if ( pSymParams )
				hResult = pDbgSymGroup->GetSymbolParameters( 0, ulCount, pSymParams );
			else
			{
				hResult = E_OUTOFMEMORY;
				pSymParams = NULL;
			}

			if ( SUCCEEDED( hResult ) )
			{
				for ( ULONG i = 0 ; i < ulCount ; i++ )
				{
					WCHAR szName[ MAX_PATH ], *pName;
					ULONG ulSize;
					DEBUG_VALUE Value;

					pName = &szName[1];
					hResult = pDbgSymGroup->GetSymbolNameWide( i, pName, MAX_PATH - 1, &ulSize ) ;
					if ( FAILED( hResult ) )
						pName[0] = 0;

					// Prefix ! so this is evaluated as symbol
					szName[0] = '!';
					hResult = g_Control->EvaluateWide( szName, DEBUG_VALUE_INVALID, &Value, NULL );
				}
			}

			if ( pSymParams )
				delete pSymParams;
		}

		pDbgSymGroup->Release();
	}*/

	return S_OK;
}

HRESULT
CDbgService::GetNearNameByOffset (
	__in ULONG64 Offset,
    __in LONG Delta,
    __out BSTR* pSymbol,
    __out ULONG64* pDisplacement
    )
{
    WCHAR Name[0x1000];
    ULONG NameSize = 0;
    HRESULT hResult = g_Symbols->GetNearNameByOffsetWide( Offset, Delta, (PWSTR) Name, _countof(Name), &NameSize, pDisplacement );
    if ( SUCCEEDED( hResult ) )
        *pSymbol = SysAllocString( Name );
    else
        *pSymbol = NULL;

    return S_OK;
}


HRESULT
CDbgService::SetSymbolPath (
    __in BSTR pSymbolPath
    )
{
    HRESULT hResult = g_Symbols->SetSymbolPathWide( pSymbolPath );
    SysFreeString( pSymbolPath );

    PWCHAR pwchReload = L".reload";
    BSTR bstr = SysAllocString( pwchReload );
    if (bstr)
    {
        Command_Push( 0, 0, bstr );
        SysFreeString( bstr );
    }

    return hResult;
}

HRESULT
CDbgService::GetCurrentProcessId (
    __out PULONG64 pId
    )
{
    if ( gbKernel )
    {
        ULONG64 Offset = 0;
        HRESULT hResult = g_SystemObjects->GetImplicitProcessDataOffset( &Offset );

        if (SUCCEEDED( hResult ) )
        {
            *pId = 0;
            ULONG TypeId;
            ULONG64 Module;
            hResult = g_Symbols->GetSymbolTypeIdWide( L"_EPROCESS", &TypeId, &Module );
        
            if ( SUCCEEDED( hResult ) )
            {
                ULONG FieldOffset = 0;
                hResult = g_Symbols->GetFieldOffsetWide( Module, TypeId, L"UniqueProcessId", &FieldOffset );

                if ( SUCCEEDED( hResult ) )
                {
                    ULONG64 Value = 0;
                    ULONG BytesRead;
                    if (S_OK == IsPointer64Bit() )
                        hResult = g_Data->ReadVirtual( Offset + FieldOffset, &Value, sizeof(Value), &BytesRead );
                    else
                        hResult = g_Data->ReadVirtual( Offset + FieldOffset, &Value, sizeof(ULONG), &BytesRead );

                    if ( SUCCEEDED( hResult ) )
                        *pId = Value;
                }
            }
        }
        
        return hResult;
    }

    *pId = 0;
    return g_SystemObjects->GetCurrentProcessId( (PULONG) pId );
}

HRESULT
CDbgService::GetCurrentThreadId (
    __out PULONG64 pId
    )
{
    if ( gbKernel )
    {
        ULONG64 Offset = 0;
        HRESULT hResult = g_SystemObjects->GetImplicitThreadDataOffset( &Offset );

        if (SUCCEEDED( hResult ) )
        {
            *pId = 0;
            ULONG TypeId;
            ULONG64 Module;
            hResult = g_Symbols->GetSymbolTypeIdWide( L"nt!_KTHREAD", &TypeId, &Module );
        
            if ( SUCCEEDED( hResult ) )
            {
                ULONG FieldOffset = 0;
                hResult = g_Symbols->GetFieldOffsetWide( Module, TypeId, L"Teb", &FieldOffset );

                if ( SUCCEEDED( hResult ) )
                {
                    ULONG64 Value = 0;
                    ULONG BytesRead;
                    if (S_OK == IsPointer64Bit() )
                        hResult = g_Data->ReadVirtual( Offset + FieldOffset, &Value, sizeof(Value), &BytesRead );
                    else
                        hResult = g_Data->ReadVirtual( Offset + FieldOffset, &Value, sizeof(ULONG), &BytesRead );

                    if ( !SUCCEEDED( hResult ) )
                        return S_OK;

                    if ( !Value )   // idle thread
                        return S_OK;

                    hResult = g_Symbols->GetSymbolTypeIdWide( L"nt!_TEB", &TypeId, &Module );
                    if ( SUCCEEDED( hResult ) )
                        hResult = g_Symbols->GetFieldOffsetWide( Module, TypeId, L"ClientId.UniqueThread", &FieldOffset );
                    
                    if ( SUCCEEDED( hResult ) )
                    {
                        if (S_OK == IsPointer64Bit() )
                            hResult = g_Data->ReadVirtual( Value + FieldOffset, &Value, sizeof(Value), &BytesRead );
                        else
                            hResult = g_Data->ReadVirtual( Value + FieldOffset, &Value, sizeof(ULONG), &BytesRead );
                    }
                    
                    if ( SUCCEEDED( hResult ) )
                        *pId = Value;
                }
            }
        }
        
        return hResult;
    }

    *pId = 0;
    return g_SystemObjects->GetCurrentThreadId( (PULONG) pId );
}